import random
import math


def euclidean_distance(point1, point2):
    """计算欧几里得距离"""
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(point1, point2)))


class KMeans:
    def __init__(self, k=3, max_iters=100):
        self.k = k
        self.max_iters = max_iters
        self.centroids = []
        self.clusters = []

    def initialize_centroids(self, data):
        """随机初始化质心"""
        random_indices = random.sample(range(len(data)), self.k)
        self.centroids = [data[i] for i in random_indices]

    def assign_clusters(self, data):
        """将数据点分配到最近的质心"""
        clusters = [[] for _ in range(self.k)]

        for point in data:
            distances = [euclidean_distance(point, centroid) for centroid in self.centroids]
            closest_centroid = distances.index(min(distances))
            clusters[closest_centroid].append(point)

        return clusters

    def update_centroids(self, clusters):
        """更新质心位置"""
        new_centroids = []
        for cluster in clusters:
            if cluster:
                # 计算簇中所有点的均值作为新质心
                new_centroid = [sum(dim) / len(cluster) for dim in zip(*cluster)]
                new_centroids.append(new_centroid)
            else:
                # 如果簇为空，保持原质心
                new_centroids.append(self.centroids[len(new_centroids)])

        return new_centroids

    def fit(self, data):
        """训练模型"""
        self.initialize_centroids(data)

        for _ in range(self.max_iters):
            self.clusters = self.assign_clusters(data)
            new_centroids = self.update_centroids(self.clusters)

            # 检查质心是否收敛
            if all(euclidean_distance(old, new) < 1e-6
                   for old, new in zip(self.centroids, new_centroids)):
                break

            self.centroids = new_centroids

    def predict(self, data):
        """预测数据点所属的簇"""
        predictions = []
        for point in data:
            distances = [euclidean_distance(point, centroid) for centroid in self.centroids]
            predictions.append(distances.index(min(distances)))
        return predictions


# 测试代码
if __name__ == "__main__":
    # 生成示例数据
    data = []
    # 簇1
    for _ in range(20):
        data.append([random.gauss(2, 0.5), random.gauss(2, 0.5)])
    # 簇2
    for _ in range(20):
        data.append([random.gauss(8, 0.5), random.gauss(8, 0.5)])
    # 簇3
    for _ in range(20):
        data.append([random.gauss(5, 0.5), random.gauss(2, 0.5)])

    kmeans = KMeans(k=3)
    kmeans.fit(data)

    print("质心位置:")
    for i, centroid in enumerate(kmeans.centroids):
        print(f"簇{i}: {centroid}")

    print(f"\n每个簇的数据点数量:")
    for i, cluster in enumerate(kmeans.clusters):
        print(f"簇{i}: {len(cluster)}个点")